# encoding: utf-8

import json
import numpy as np
from typing import List
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn

from zcbot_web_core.model.handler import Resp

from sku_extract import SkuExtractor

pretrained_model = r'E:/work/models/info_extract/model_189000'
extractor = SkuExtractor(pretrained_model)

app = FastAPI()


# Extend the JSONEncoder class
class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.floating):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)


class Item(BaseModel):
    text: str = None
    textList: List[str] = None


def ner_one(text: str):
    result = extractor.extract(text)
    return result


def ner_batch(text_list: List[str]):
    result = extractor.extract_batch(text_list)
    return result


@app.get("/uie-ner")
def get_ner(text: str):
    result = ner_one(text)
    return Resp.ok(result)


@app.post("/uie-ner")
def post_ner(item: Item):
    text_list = []
    if item.text:
        text_list.append(item.text)
        # result = ner_one(item.text)
    elif item.textList:
        text_list.extend(item.textList)
        # result = ner_batch(item.textList)
    if text_list:
        result = ner_batch(text_list)
    else:
        return Resp.err("body cannot be null")
    return Resp.ok(result)


if __name__ == '__main__':
    uvicorn.run(app, host="0.0.0.0", port=6697)
